#include <cusp/array1d.h>
#include <cusp/array2d.h>
#include <cusp/multiply.h>
#include <cusp/transpose.h>
#include <cusp/io/matrix_market.h>
#include <cusp/qr.h>
#include <cusp/device/block/qr.h>
#include <cusp/qrhh.h>
#include <cusp/qrhhgemv.h>
#include <cusp/qrhhb.h>

#include <iostream>
#include "../timer.h"
#include <cusp/print.h>

// #include <magma.h>

typedef float  ValueType;
typedef cusp::array2d<ValueType,cusp::host_memory>   HostMatrix;
typedef cusp::array2d<ValueType,cusp::device_memory> DeviceMatrix;
typedef cusp::array2d<ValueType,cusp::device_memory,cusp::column_major> DeviceMatrixColMaj;


int check_result(HostMatrix& Ahost, const DeviceMatrix& Q, const DeviceMatrix& R)
{
	int i = 0, f = 0;

	HostMatrix Qhost = Q;
	HostMatrix Rhost = R;
	HostMatrix Bhost;

	cusp::multiply(Qhost, Rhost, Bhost);

	if (Ahost.num_rows != Bhost.num_rows || Ahost.num_cols != Bhost.num_cols)
	{
		std::cerr << "Check FAILED! (size)" << std::endl;
		return -1;
	}

	typedef HostMatrix::values_array_type::iterator Iterator;
	Iterator first1 = Ahost.values.begin();
	Iterator last1  = Ahost.values.end();
	Iterator first2 = Bhost.values.begin();
	while(first1 != last1)
	{
		/* same condition used in unit testing */
		if (abs(*first1 - *first2) > 1e-4 * (abs(*first1) + abs(*first2)) + 1e-4)
		{
			++f;
		}
		first1++;
		first2++;
		++i;
	}

	if (f != 0)
		std::cerr << "Check FAILED! (" << f << ")" << std::endl;

	return f;
}

void print_time(const char* name, const float time)
{
	//std::cout << name << ":\t" << time << " s" << std::endl;
	std::cerr << name << ":\t" << time << " s" << std::endl;
	std::cout << "\t" << time;
}

template<typename RandomAccessIterator1, typename RandomAccessIterator2>
__global__
void qr_block_kernel_row(RandomAccessIterator1 A, size_t num_rows, size_t num_cols, size_t row_number, RandomAccessIterator2 Q)
{
	if (blockIdx.x == 0)
		cusp::device::block::qr(A, num_rows, num_cols, row_number, Q);
}

template<typename RandomAccessIterator1, typename RandomAccessIterator2>
__global__
void qr_block_kernel(RandomAccessIterator1 A, size_t num_rows, size_t num_cols, RandomAccessIterator2 Q)
{
	if (blockIdx.x == 0)
		cusp::device::block::qr(A, num_rows, num_cols, Q);
}

int main(int argc, char** argv)
{

	if (argc < 2)
	{
		std::cerr << "Usage: test <matrix.mtx>" << std::endl;
		return 1;
	}

	ValueType *tau;
	HostMatrix Ahost;
	std::cerr << "Reading matrix from file: " << argv[1] << std::endl;
	cusp::io::read_matrix_market_file(Ahost, std::string(argv[1]));
	std::cerr << "Size " << Ahost.num_rows << " x " << Ahost.num_cols << std::endl;

        double m = (double)Ahost.num_rows;
        double n = (double)Ahost.num_cols;
        double mflops = 4.0 * n * ( m * (m - n) + n*n/3.0) / 1000000.0;
        tau = (ValueType*)malloc(min(m,n)*sizeof(ValueType));
        int info;

	DeviceMatrix A = Ahost;
	DeviceMatrix Q, R;


#if 0
	{
	timer t;
	cusp::qr(A, Q, R);
	float time = t.seconds_elapsed();
	// print_time("qr_givens", time);
	std::cout << "qr_givens (cusp::blas): " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
	check_result(Ahost, Q, R);
	}
#endif

	{
	cusp::array1d<ValueType,cusp::device_memory> q4(A.num_rows, 0);
	timer t;
	qr_block_kernel_row<<<1, 256>>>(A.values.begin(), A.num_rows, A.num_cols, 3, q4.begin());
	float time = t.seconds_elapsed();
 	// print_time("qr_givens_block", time);
        std::cout << "qr_givens_block_one_row: " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
        cusp::print(q4);

	// check_result(Ahost, Q, A);
	}


	{
	cusp::array2d<ValueType,cusp::device_memory,cusp::column_major> Qt(A.num_rows, A.num_rows, 0.0);
	timer t;
	qr_block_kernel<<<1, 256>>>(A.values.begin(), A.num_rows, A.num_cols, Qt.values.begin());
	float time = t.seconds_elapsed();
 	// print_time("qr_givens_block", time);
        std::cout << "qr_givens_block: " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;

	Q = Qt;
	check_result(Ahost, Q, A);
	}

#if 0
	{
	DeviceMatrixColMaj Ac = Ahost;
	DeviceMatrixColMaj Qc, Rc;
	timer t;
	cusp::qrhh(Ac, Qc, Rc);
	float time = t.seconds_elapsed();
	// print_time("qrhh", time);
	std::cout << "qr householder (cusp): " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
	Q = Qc;
	R = Rc;
	check_result(Ahost, Q, R);
	}
#endif

	{
	DeviceMatrixColMaj Ac = Ahost;
	DeviceMatrixColMaj Qc, Rc;
	timer t;
	cusp::qrhhgemv(Ac, Qc, Rc);
	float time = t.seconds_elapsed();
	// print_time("qrhhgemv", time);
	std::cout << "qr householder (cublas): " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
	Q = Qc;
	R = Rc;
	check_result(Ahost, Q, R);
	}

#if 0
	{
	DeviceMatrixColMaj Ac = Ahost;
	DeviceMatrixColMaj Qc, Rc;
	timer t;
	magma_sgeqrf2_gpu( m, n, thrust::raw_pointer_cast(&Ac.values[0]), m, tau, &info);
	float time = t.seconds_elapsed();
	// print_time("qrhhgemv", time);
	std::cout << "MAGMA qr: " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
	Q = Qc;
	R = Rc;
	check_result(Ahost, Q, R);
	}

	{
	DeviceMatrixColMaj Ac = Ahost;
	DeviceMatrixColMaj Qc, Rc;
	timer t;
	cusp::qrhhb(Ac, Qc, Rc, 3);
	float time = t.seconds_elapsed();
	// print_time("qrhhb", time);
	std::cout << "qr block householder (cublas): " << time << " s " << mflops/((double)time) << " mflop/s " << std::endl;
	Q = Qc;
	R = Rc;
	check_result(Ahost, Q, R);
	}
#endif

	return 0;
}

